package com.cdn.clock.component.lock;

import com.cdn.clock.component.LockComponent;
import com.cdn.clock.config.ClogThreadPoolTask;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import javax.annotation.Resource;
import java.util.Collections;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

/**
 * @author 蔡定努
 * 2024/06/04 10:01
 */
public class RedisLockComponent implements LockComponent {

    @Resource(name = "redisLockTemplate")
    private RedisTemplate<String, Object> redisTemplate;
    @Resource
    private ClogThreadPoolTask clogThreadPoolTask;


    private final ConcurrentHashMap<String, ScheduledFuture<?>> renewalTasks = new ConcurrentHashMap<>();
    private final Logger logger = LoggerFactory.getLogger(RedisLockComponent.class);


    /**
     * 释放锁
     *
     * @param lockKey key
     * @return boolean
     */
    @Override
    public boolean unLock(String lockKey) {
        stopRenewalTask(lockKey);
        String RELEASE_LOCK_SCRIPT =
                "if redis.call('GET', KEYS[1]) == ARGV[1] then\n" +
                        "    return redis.call('DEL', KEYS[1])\n" +
                        "else\n" +
                        "    return 0\n" +
                        "end";
        RedisScript<Long> script = new DefaultRedisScript<>(RELEASE_LOCK_SCRIPT, Long.class);
        Long result = redisTemplate.execute(script, Collections.singletonList(lockKey), 1);
        LoggerFactory.getLogger(RedisLockComponent.class).info("unlock {},result is {}", lockKey, result);
        return result != null && result == 1;
    }


    @Override
    public boolean lock(String lockKey, long waitTime, long leaseTime, TimeUnit timeUnit) {
        String LOCK_SCRIPT = "local expireTime = tonumber(ARGV[1]) " +
                "if redis.call('SETNX', KEYS[1], 1) == 1 then " +
                "   if expireTime > 0 then " +
                "       redis.call('PEXPIRE', KEYS[1], expireTime) " +
                "   end " +
                "   return 1 " +
                "else " +
                "   return 0 " +
                "end";

        long millis = timeUnit.toMillis(leaseTime);
        RedisScript<Long> script = new DefaultRedisScript<>(LOCK_SCRIPT, Long.class);
        Long result = redisTemplate.execute(script, Collections.singletonList(lockKey), millis);
        boolean success = result != null && result == 1L;
        if (waitTime > 0 && !success) {
            long end = System.currentTimeMillis() + timeUnit.toMillis(waitTime);
            while (System.currentTimeMillis() < end) {
                try {
                    Thread.sleep(100);
                    result = redisTemplate.execute(script, Collections.singletonList(lockKey), millis);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    return false;
                }
            }
        }
        success = result != null && result == 1L;
        if (success) {
            startRenewalTask(lockKey, millis);
        }
        return success;
    }


    /**
     * 续期
     *
     * @param lockKey
     * @param leaseMillis
     */
    private void startRenewalTask(String lockKey, long leaseMillis) {
        stopRenewalTask(lockKey); // Stop any existing renewal task for the same lock
        if (leaseMillis > 0) {
            ScheduledFuture<?> renewalTask = clogThreadPoolTask.getTaskScheduler().scheduleAtFixedRate(() -> {
                String RENEW_SCRIPT = "if redis.call('GET', KEYS[1]) == '1' then " +
                        "redis.call('PEXPIRE', KEYS[1], tonumber(ARGV[1])) " +
                        "return 1 " +
                        "else " +
                        "return 0 " +
                        "end";
                RedisScript<Long> script = new DefaultRedisScript<>(RENEW_SCRIPT, Long.class);
                Long execute = redisTemplate.execute(script, Collections.singletonList(lockKey), leaseMillis);
                if (execute == null || execute == 0) {
                    stopRenewalTask(lockKey);
                    return;
                }

                logger.info("Renewed lock: {},expire at {}", lockKey, leaseMillis);
            }, leaseMillis / 2); // Renew at half the lease time
            renewalTasks.put(lockKey, renewalTask);
        }
    }


    private void stopRenewalTask(String lockKey) {
        ScheduledFuture<?> renewalTask = renewalTasks.get(lockKey);
        if (renewalTask != null) {
            renewalTask.cancel(false);
        }
        logger.info("Stop renewal task for lock: {}", lockKey);
    }


    @Override
    public boolean forceUnlock(String lockKey) {
        stopRenewalTask(lockKey);
        return unLock(lockKey);
    }

    @Override
    public boolean lock(String lockKey, long leaseTime, TimeUnit timeUnit) {
        return lock(lockKey, -1, leaseTime, timeUnit);
    }

    @Override
    public boolean lock(String lockKey) {
        return lock(lockKey, -1L, -1, TimeUnit.MILLISECONDS);
    }

}
